#!/usr/bin/env python

import json
import imaplib
import email
import email.encoders
from email.mime.base import MIMEBase
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
import mimetypes
import subprocess
import time
import platform
import os
import zipfile
import shutil
import io
import encrypted_packet_type_b_pynacl as packet_type_b


class IMAPDraftBox:
	def __init__(self, name, config):
		self.name = name
		self.config = config


	def get_payload(self, boxname="input_box"):
		ret = None

		if self.config["imap_ssl"]:
			M = imaplib.IMAP4_SSL(self.config["imap_server"], self.config["imap_port"])
		else:
			M = imaplib.IMAP4(self.config["imap_server"], self.config["imap_port"])
		M.login(self.config["username"], self.config["password"])
		M.select(self.config[boxname] if boxname == "input_box" or boxname == "output_box" else boxname)

		typ, data = M.search(None, 'ALL')
		for num in data[0].split():
			typ, data = M.fetch(num, '(RFC822)')
			M.store(num, '+FLAGS', '\\Deleted')
			msg = email.message_from_string(data[0][1])
			print("mail %s: %s" % (num, msg["Subject"]))
			for part in msg.walk():
				if part.get_content_maintype() == 'multipart':
					continue
				filename = part.get_filename()
				if not filename is None and filename.endswith(".bin"):
					ret = part.get_payload(decode=True)
					break
			if not ret is None:
				print("Found payload")
				break

		M.expunge()
		M.close()
		M.logout()

		return ret


	def upload_payload(self, payload, boxname="output_box", subject="Encrypted response"):
		msg = MIMEMultipart()
		msg["From"] = self.config["username"]
		msg["To"] = self.config["username"]
		msg["Subject"] = subject
		msg.attach(MIMEText("See the attachments...", "plain"))

		part = MIMEBase("application", "octet-stream")
		part.set_payload(payload)
		email.encoders.encode_base64(part)
		part.add_header(
			"Content-Disposition",
			"attachment; filename=payload.bin",
		)

		msg.attach(part)

		if self.config["imap_ssl"]:
			M = imaplib.IMAP4_SSL(self.config["imap_server"], self.config["imap_port"])
		else:
			M = imaplib.IMAP4(self.config["imap_server"], self.config["imap_port"])
		M.login(self.config["username"], self.config["password"])
		M.append(self.config[boxname] if boxname == "input_box" or boxname == "output_box" else boxname,
			'', imaplib.Time2Internaldate(time.time()), msg.as_string()
		)
		M.logout()


def save_payload(payload, filename, default_ext=None):
	m = ord(payload[0]) & 0xFF
	ext = payload[1:m+1]
	payload = payload[m+1:]

	if m > 0:
		filename += "." + ext
	elif not default_ext is None:
		filename += "." + default_ext

	with open(filename, 'wb') as fp:
		fp.write(payload)

	return filename


def run_payload(config, payload):
	ret = b''

	m = ord(payload[0]) & 0xFF
	ext = payload[1:m+1]
	payload = payload[m+1:]

	filename = None
	is_zip_file = ext == b'zip'

	if is_zip_file:
		try:
			zipfile.ZipFile(io.BytesIO(payload)).extractall("build/payload")
		except Exception as e:
			ret += b'Failed to extract the zip archive: %s\n' % e
			return ret
		filename = None
		for fn in os.listdir("build/payload"):
			if fn.startswith("payload"):
				filename = "build/payload/" + fn
				break
	else:
		filename = "build/payload"
		if m > 0:
			filename += "." + ext
		try:
			with open(filename, 'wb') as fp:
				fp.write(payload)
		except Exception as e:
			ret += b'Failed to create file %s: %s\n' % (filename, e)
			return ret

	if not filename is None:
		if platform.system() == "Windows":
			filename = filename.replace("/", "\\")
		else:
			filename = filename.replace("\\", "/")
			os.chmod(filename, 0755)

		ss = time.time()
		try:
			env = dict(os.environ)
			try:
				env.update(config["env"])
			except:
				pass
			ret += subprocess.check_output(filename, stderr=subprocess.STDOUT, env=env)
			ret += b'\n\nProgram exits normally\n'
		except subprocess.CalledProcessError as e:
			ret += e.output
			ret += b'\n\nProgram returns %d\n' % e.returncode
		except Exception as e:
			ret += b'Failed to run program %s: %s\n' % (filename, e)
		ss = time.time() - ss
		mm = int(ss / 60.0)
		ss -= mm * 60.0
		hh = mm / 60
		mm -= hh * 60
		ret += b'Running time: %02d:%02d:%02.3f\n' % (hh, mm, ss)
	else:
		ret += b'No payload found in zip archive\n'

	try:
		if is_zip_file:
			shutil.rmtree("build/payload")
		else:
			os.remove(filename)
	except Exception as e:
		ret += b'Error when deleting payload: %s\n' % e

	return ret


def create_payload(payload, ext, keys):
	m = bytes(bytearray([len(ext) & 0xFF]))
	return packet_type_b.create_encrypted_packet(keys, m + ext + payload)


def create_inputs(config):
	inputs = config["inputs"]
	ret = []
	for name in config["used_inputs"]:
		input = inputs[name]
		type = input["type"]
		if type == "imap_draftbox":
			ret.append(IMAPDraftBox(name, input))
		else:
			print("Unknown input type: %s" % type)
	return ret


def main():
	# load config files
	config = json.load(open("build/config.json", 'r'))
	keys = json.load(open("build/server_key.json", 'r'))

	inputs = create_inputs(config)

	for input in inputs:
		print("Checking %s" % input.name)
		while True:
			payload = input.get_payload()
			if payload is None:
				break
			try:
				payload = packet_type_b.verify_encrypted_packet(keys, payload)
			except Exception as e:
				print "Failed to decrypt message: ", e
			else:
				ret = run_payload(config, payload)
				print ret
				print "Sending the reply..."
				input.upload_payload(create_payload(ret, b'txt', keys))


if __name__ == '__main__':
	main()
